import os
import itertools
import datetime
import editdistance
import numpy as np
import pylab
from keras import backend as K
from keras.layers.convolutional import Conv2D, MaxPooling2D
from keras.layers import Input, Dense, Activation
from keras.layers import Reshape, Lambda
from keras.layers.merge import add, concatenate
from keras.models import Model
from keras.layers.recurrent import GRU
from keras.optimizers import SGD
from keras.utils import plot_model
import keras.callbacks
from dataset import DatasetGenerator

OUTPUT_DIR = 'output'
n_channel = 3
np.random.seed(55)


# the actual loss calc occurs here despite it not being
# an internal Keras loss function

def ctc_lambda_func(args):
    y_pred, labels, input_length, label_length = args
    # the 2 is critical here since the first couple outputs of the RNN
    # tend to be garbage:
    y_pred = y_pred[:, 2:, :]
    return K.ctc_batch_cost(labels, y_pred, input_length, label_length)


class VizCallback(keras.callbacks.Callback):
    def __init__(self, run_name, test_func, text_img_gen, test_img_gen=None, num_display_words=6):
        self.test_func = test_func
        self.run_name = run_name
        self.output_dir = os.path.join(OUTPUT_DIR, run_name)
        self.text_img_gen = text_img_gen
        self.test_img_gen = test_img_gen
        self.num_display_words = num_display_words
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)

    # For a real OCR application, this should be beam search with a dictionary
    # and language model.  For this example, best path is sufficient.
    def decode_batch(self, test_func, word_batch):
        # out's shape:(minibatch, timesteps, classes)
        out = test_func([word_batch])[0]
        ret = []
        for j in range(out.shape[0]):
            # out_best's shape:(timesteps - 2, 1)
            out_best = list(np.argmax(out[j, 2:], 1))
            # print 'out_best:{},{}'.format(out_best, len(out_best))
            out_best = [k for k, g in itertools.groupby(out_best)]
            # print 'out_best_after_groupby:{},{}'.format(out_best, len(out_best))
            # 26 is space, 27 is CTC blank char
            outstr = ''
            for c in out_best:
                if 0 <= c < (self.text_img_gen.num_classes - 2):
                    outstr += self.text_img_gen.classes[c]
                elif c == (self.text_img_gen.num_classes - 2):
                    outstr += ' '
                # else:
                    # print "blank!",
            ret.append(outstr)
        # print "decode:{}".format(ret)

        return ret

    def show_edit_distance(self, num):
        num_left = num
        mean_norm_ed = 0.0
        mean_ed = 0.0
        while num_left > 0:
            word_batch = next(self.text_img_gen.next_batch())[0]
            num_proc = min(word_batch['the_input'].shape[0], num_left)
            print "num_left:", num_left
            decoded_res = self.decode_batch(self.test_func, word_batch['the_input'][0:num_proc])
            for j in range(0, num_proc):
                print "show:{}, {},length of decode:{}".format(word_batch['source_str'][j].encode('utf-8'), decoded_res[j].encode('utf-8'),
                                                               len(decoded_res[j]))
                edit_dist = editdistance.eval(decoded_res[j], word_batch['source_str'][j])
                mean_ed += float(edit_dist)
                mean_norm_ed += float(edit_dist) / len(word_batch['source_str'][j])
            num_left -= num_proc
        mean_norm_ed = mean_norm_ed / num
        mean_ed = mean_ed / num
        print('\nOut of %d samples:  Mean edit distance: %.3f Mean normalized edit distance: %0.3f'
              % (num, mean_ed, mean_norm_ed))

    def on_epoch_end(self, epoch, logs={}):

        print "start to save weights"
        self.model.save_weights(os.path.join(self.output_dir, 'weights%02d.h5' % (epoch)))
        print "start to show edit distance"
        self.show_edit_distance(256)
        print "show decode_batch"
        word_batch = next(self.text_img_gen.next_batch())[0]
        res = self.decode_batch(self.test_func, word_batch['the_input'][0:self.num_display_words])
        cols = 2
        for i in range(self.num_display_words):
            pylab.subplot(self.num_display_words // cols, cols, i + 1)
            if K.image_data_format() == 'channels_first':
                the_input = word_batch['the_input'][i, :, :, :]
                gray_img = the_input.transpose((1, 0, 2))[:, :, 0]
                pylab.imshow(gray_img)
            else:
                the_input = word_batch['the_input'][i, :, :, :]
                pylab.imshow(the_input.transpose((1, 0, 2)))
            pylab.xlabel('Truth = \'%s\'\nDecoded = \'%s\'' % (word_batch['source_str'][i], res[i]))
            print "{} _ {}".format(epoch, i)
            print "Source_str:{}".format(word_batch['source_str'][i].encode('utf-8'))
            print "Decode_str:{}".format(res[i])
        fig = pylab.gcf()
        fig.set_size_inches(10, 13)
        pylab.savefig(os.path.join(self.output_dir, 'e%02d.png' % (epoch)))
        pylab.close()

    def on_train_begin(self, logs=None):
        print "start to save model structure"
        json_string = self.model.to_json()
        open(os.path.join(self.output_dir, 'model_structure.json'), 'w').write(json_string)
        print "start to plot model structure"
        plot_model(self.model, to_file=os.path.join(self.output_dir, 'model.png'), show_shapes=True, show_layer_names=True)

    def on_train_end(self, logs=None):
        print "start to save model"
        self.model.save(os.path.join(self.output_dir, 'model.h5'))
        if self.test_img_gen is not None:
            output_path = os.path.join(self.test_img_gen.datapath, 'results')
            if not os.path.exists(output_path):
                os.mkdir(output_path)
            output = os.path.join(output_path, self.run_name + '.txt')
            num_test_images = self.test_img_gen.num_images
            out_lines = []
            for i in range(num_test_images):
                im_name = self.test_img_gen.images_index[i] + '.' + self.test_img_gen.image_ext
                im = self.test_img_gen.get_image(i)
                res = self.decode_batch(self.test_func, im[np.newaxis, :, :, :])
                # print "test_res:", res, type(res)
                out_line = '{}, "{}"'.format(im_name, res[0].encode('utf-8'))
                out_lines.append(out_line)
                # print 'source:{}, {}'.format(self.test_img_gen.labels_index[i].encode('utf-8'), out_line)
            out = '\r\n'.join(out_lines)
            with open(output, 'w') as f:
                f.write(out)


def train(run_name, start_epoch, stop_epoch, img_w):
    # Input Parameters
    img_h = 64

    # Network parameters
    conv_filters = 16
    kernel_size = (3, 3)
    pool_size = 2
    time_dense_size = 50
    rnn_size = 512
    train_minibatch_size = 128
    val_minibatch_size = 64

    if K.image_data_format() == 'channels_first':
        input_shape = (n_channel, img_w, img_h)
    else:
        input_shape = (img_w, img_h, n_channel)

    datapath = '/home/hsp/jr/work/ocr/data'
    dataset = 'COCO-Text'
    train_img_gen = DatasetGenerator(datapath=datapath,
                                     dataset=dataset,
                                     image_set='train',
                                     minibatch_size=train_minibatch_size,
                                     img_w=img_w,
                                     img_h=img_h,
                                     downsample_factor=(pool_size ** 2),
                                     image_ext='jpg',
                                     absolute_max_string_len=50)
    print "load train dataset done!"
    val_img_gen = DatasetGenerator(datapath=datapath,
                                   dataset=dataset,
                                   image_set='val',
                                   minibatch_size=val_minibatch_size,
                                   img_w=img_w,
                                   img_h=img_h,
                                   downsample_factor=(pool_size ** 2),
                                   image_ext='jpg',
                                   absolute_max_string_len=50)
    print "load val dataset done!"
    test_img_gen = DatasetGenerator(datapath=datapath,
                                    dataset=dataset,
                                    image_set='test',
                                    minibatch_size=1,
                                    img_w=img_w,
                                    img_h=img_h,
                                    downsample_factor=(pool_size ** 2),
                                    image_ext='jpg',
                                    absolute_max_string_len=50)
    print "load test dataset done!"

    '--------------------------------------- model -----------------------------------------'
    act = 'relu'
    input_data = Input(name='the_input', shape=input_shape, dtype='float32')
    inner = Conv2D(conv_filters, kernel_size, padding='same',
                   activation=act, kernel_initializer='he_normal',
                   name='conv1')(input_data)
    inner = MaxPooling2D(pool_size=(pool_size, pool_size), name='max1')(inner)
    inner = Conv2D(conv_filters, kernel_size, padding='same',
                   activation=act, kernel_initializer='he_normal',
                   name='conv2')(inner)
    inner = MaxPooling2D(pool_size=(pool_size, pool_size), name='max2')(inner)

    conv_to_rnn_dims = (img_w // (pool_size ** 2), (img_h // (pool_size ** 2)) * conv_filters)
    inner = Reshape(target_shape=conv_to_rnn_dims, name='reshape')(inner)

    # cuts down input size going into RNN:
    inner = Dense(time_dense_size, activation=act, name='dense1')(inner)

    # Two layers of bidirectional GRUs
    # GRU seems to work as well, if not better than LSTM:
    gru_1 = GRU(rnn_size, return_sequences=True, kernel_initializer='he_normal', name='gru1')(inner)
    gru_1b = GRU(rnn_size, return_sequences=True, go_backwards=True, kernel_initializer='he_normal', name='gru1_b')(
        inner)
    gru1_merged = add([gru_1, gru_1b])
    gru_2 = GRU(rnn_size, return_sequences=True, kernel_initializer='he_normal', name='gru2')(gru1_merged)
    gru_2b = GRU(rnn_size, return_sequences=True, go_backwards=True, kernel_initializer='he_normal', name='gru2_b')(
        gru1_merged)

    # transforms RNN output to character activations:
    inner = Dense(train_img_gen.num_classes, kernel_initializer='he_normal',
                  name='dense2')(concatenate([gru_2, gru_2b]))
    y_pred = Activation('softmax', name='softmax')(inner)
    Model(inputs=input_data, outputs=y_pred).summary()

    labels = Input(name='the_labels', shape=[train_img_gen.absolute_max_string_len], dtype='float32')
    input_length = Input(name='input_length', shape=[1], dtype='int64')
    label_length = Input(name='label_length', shape=[1], dtype='int64')
    # Keras doesn't currently support loss funcs with extra parameters
    # so CTC loss is implemented in a lambda layer
    loss_out = Lambda(ctc_lambda_func, output_shape=(1,), name='ctc')([y_pred, labels, input_length, label_length])

    # clipnorm seems to speeds up convergence
    sgd = SGD(lr=0.02, decay=1e-6, momentum=0.9, nesterov=True, clipnorm=5)

    model = Model(inputs=[input_data, labels, input_length, label_length], outputs=loss_out)

    # the loss calc occurs elsewhere, so use a dummy lambda func for the loss
    model.compile(loss={'ctc': lambda y_true, y_pred: y_pred}, optimizer=sgd)
    if start_epoch > 0:
        weight_file = os.path.join(OUTPUT_DIR, os.path.join(run_name, 'weights%02d.h5' % (start_epoch - 1)))
        model.load_weights(weight_file)
    # captures output of softmax so we can decode the output during visualization
    test_func = K.function([input_data], [y_pred])

    viz_cb = VizCallback(run_name, test_func, val_img_gen, test_img_gen, num_display_words=8)

    steps_per_epoch = train_img_gen.num_images // train_minibatch_size
    validation_steps = val_img_gen.num_images // val_minibatch_size
    print 'steps_per_epoch:{}'.format(steps_per_epoch)
    print 'validation_steps:{}'.format(validation_steps)

    model.fit_generator(generator=train_img_gen.next_batch(),
                        steps_per_epoch=steps_per_epoch,
                        epochs=stop_epoch,
                        validation_data=val_img_gen.next_batch(),
                        validation_steps=validation_steps,
                        callbacks=[viz_cb, val_img_gen],
                        initial_epoch=start_epoch)


if __name__ == '__main__':
    run_name = datetime.datetime.now().strftime('%Y:%m:%d:%H:%M:%S')
    train('COCO-Text_time_50_test_' + run_name, 0, 50, 256)
    # increase to wider images and start at epoch 20. The learned weights are reloaded
    # train('dataset_time_20_' + run_name, 100, 125, 192)
